#include "stdafx.h"
#include "Robot.h"
#include "OgreUtils.h"
using namespace std;

Robot::Robot(World* world):
world(world)
{
	createBody();
	createNet();
}


void Robot::update()
{
	vector<float> in;
	in.resize(objects.size()*13);
	for(size_t i = 0; i < objects.size(); i++)
	{
		const btVector3& pos = objects[i]->body->getCenterOfMassPosition();
		btQuaternion rot = objects[i]->body->getOrientation();
		const btVector3& speed = objects[i]->body->getLinearVelocity();
		const btVector3& angSpeed = objects[i]->body->getAngularVelocity();
		in[i*13 + 0] = pos.x();
		in[i*13 + 1] = pos.y();
		in[i*13 + 2] = pos.z();
		in[i*13 + 3] = rot.x();
		in[i*13 + 4] = rot.y();
		in[i*13 + 5] = rot.z();
		in[i*13 + 6] = rot.w();
		in[i*13 + 7] = speed.x();
		in[i*13 + 8] = speed.y();
		in[i*13 + 9] = speed.z();
		in[i*13 + 10] = angSpeed.x();
		in[i*13 + 11] = angSpeed.y();
		in[i*13 + 12] = angSpeed.z();
	}

	float* out = net->run(&in[0]);
	
	int outIndex = 0;
	for(size_t i = 0; i < joints1Dof.size(); i++)
	{
		updateJoint(joints1Dof[i], 0, out[outIndex]);
		outIndex++;
	}
	for(size_t i = 0; i < joints3Dof.size(); i++)
	{
		updateJoint(joints3Dof[i], 0, out[outIndex + 0]);
		updateJoint(joints3Dof[i], 1, out[outIndex + 1]);
		outIndex += 2;
	}
}

void Robot::updateJoint(btGeneric6DofConstraint* joint, int index, float v)
{
	static const float maxVelocity = 10.f;
	if(v < -1.f) v = -1.f;
	if(v > 1.f) v = 1.f;
	v = v * maxVelocity;
	btRotationalLimitMotor* motor = joint->getRotationalLimitMotor(index);
	motor->m_enableMotor = true;
	motor->m_targetVelocity = v;
}


float* Robot::getNeuralWeights()
{
	return net->getWeights();
}

size_t Robot::getNeuralWeightsCount() const
{
	return net->getWeightsCount();
}

static GameObject* createLeg(World* world, const btTransform& trans)
{
	Ogre::SceneManager* sm = world->getSceneManager();
	Ogre::SceneNode* node = sm->getRootSceneNode()->createChildSceneNode();
	Ogre::SceneNode* boxNode1 = 0;
	Ogre::SceneNode* boxNode2 = 0;
	float mass = 0.f;
	btCollisionShape* boxShape1 = 0;
	{
		btVector3 size(2.f, 6.f, 2.f);
		mass += size.x() * size.y() * size.z();
		boxShape1 = new btBoxShape(size);
		boxNode1 = node->createChildSceneNode();
		boxNode1->attachObject(sm->createEntity("box.mesh"));
		boxNode1->setScale(size.x(), size.y(), size.z());
	}

	btCollisionShape* boxShape2 = 0;
	{
		btVector3 size(3.f, 0.5f, 5.f);
		mass += size.x() * size.y() * size.z();
		boxShape2 = new btBoxShape(size);
		boxNode2 = node->createChildSceneNode();
		boxNode2->attachObject(sm->createEntity("box.mesh"));
		boxNode2->setScale(size.x(), size.y(), size.z());
	}

	btCompoundShape* shape = new btCompoundShape();
	{
		btVector3 pos(0.f, 0.f, 0.f);
		btTransform trans;
		trans.setIdentity();
		trans.setOrigin(pos);
		shape->addChildShape(trans, boxShape1);
		boxNode1->setPosition(pos.x(), pos.y(), pos.z());
	}
	{
		btVector3 pos(0.f, -6.f, 2.f);
		btTransform trans;
		trans.setIdentity();
		trans.setOrigin(pos);
		shape->addChildShape(trans, boxShape2);
		boxNode2->setPosition(pos.x(), pos.y(), pos.z());
	}
	GameObject* obj = world->addObject(node, shape, trans, mass, false);	
	return obj;
}

Leg Robot::createLeg(GameObject* body, const btVector3& offset)
{
	Leg leg;
	btVector3 pos0 = body->body->getCenterOfMassPosition();
	btVector3 size(2.f, 6.f, 2.f);
	btVector3 pos = pos0 + offset;
	{
		btTransform trans;
		trans.setIdentity();
		trans.setOrigin(btVector3(pos.x(), pos.y() - size.y(), pos.z()));
		leg.obj0 = createBox(world, size, trans);

		trans.setOrigin(btVector3(pos.x(), pos.y() - 3*size.y(), pos.z()));
		leg.obj1 = ::createLeg(world, trans);
	}


	{
		btTransform frame1;
		frame1.setIdentity();
		frame1.setOrigin(offset);
		btTransform frame2;
		frame2.setIdentity();
		frame2.setOrigin(btVector3(0.f, size.y(), 0.f));
		leg.joint0 = new btGeneric6DofConstraint(*body->body, *leg.obj0->body, frame1, frame2, true);
		leg.joint0->setLimit(3, -SIMD_PI*0.2f, SIMD_PI*0.5f);
		leg.joint0->setLimit(4, -SIMD_PI*0.2f, SIMD_PI*0.2f);
		leg.joint0->setLimit(5, 0.f, 0.f);
		leg.joint0->getRotationalLimitMotor(0)->m_maxMotorForce = 1.f*18000.f;
		leg.joint0->getRotationalLimitMotor(1)->m_maxMotorForce = 1.f*18000.f;
		world->addJoint(leg.joint0, true);
	}
	
	{
		btTransform frame1;
		frame1.setIdentity();
		frame1.setOrigin(btVector3(0.f, -size.y(), 0.f));
		btTransform frame2;
		frame2.setIdentity();
		frame2.setOrigin(btVector3(0.f, size.y(), 0.f));
		leg.joint1 = new btGeneric6DofConstraint(*leg.obj0->body, *leg.obj1->body, frame1, frame2, true);
		leg.joint1->setLimit(3, -SIMD_PI*0.5f, SIMD_PI*0.0f);
		leg.joint1->setLimit(4, 0.f, 0.f);
		leg.joint1->setLimit(5, 0.f, 0.f);
		leg.joint1->getRotationalLimitMotor(0)->m_maxMotorForce = 1.f*7500.f;
		world->addJoint(leg.joint1, true);
	}
	
	objects.push_back(leg.obj0);
	objects.push_back(leg.obj1);
	joints3Dof.push_back(leg.joint0);
	joints1Dof.push_back(leg.joint1);
	return leg;
}

void Robot::createWheel(GameObject* body, const btVector3& offset)
{
	btVector3 pos0 = body->body->getCenterOfMassPosition();
	btVector3 size(3.f, 6.f, 6.f);
	btVector3 pos = pos0 + offset;
	GameObject* obj;
	{
		btTransform trans;
		trans.setIdentity();
		trans.setOrigin(pos);
		obj = createBox(world, size, trans);
		objects.push_back(obj);
	}
	{
		btTransform frame1;
		frame1.setIdentity();
		frame1.setOrigin(offset);
		btTransform frame2;
		frame2.setIdentity();
		frame2.setOrigin(btVector3(0.f, 0.f, 0.f));
		btGeneric6DofConstraint* joint = new btGeneric6DofConstraint(*body->body, *obj->body, frame1, frame2, true);
		joint->setLimit(4, 0.f, 0.f);
		joint->setLimit(5, 0.f, 0.f);
		world->addJoint(joint, true);
		joints1Dof.push_back(joint);
	}
}

void Robot::createBody()
{
	const float w = 6.f;	
	btTransform trans;
	trans.setIdentity();
	{
		trans.setOrigin(btVector3(0.f, 37.f, 0.f));
		btVector3 size(8.f, 10.f, 2.f);
		body = createBox(world, size, trans);
		objects.push_back(body);
	}
	
	createLeg(body, btVector3(w, -10.f, 0.f));
	createLeg(body, btVector3(-w, -10.f, 0.f));
}

void Robot::createNet()
{
	size_t nIn = objects.size() * 13;
	vector<size_t> layers;
	layers.push_back(100);
	layers.push_back(joints1Dof.size() + joints3Dof.size() * 2);
	net.reset(new NeuralNet(nIn, layers));
}

